import os
import json
import numpy as np
import matplotlib.colors as mcolors

from PIL import Image
from .structure import Point3D
from typing import Any, Dict, List, Tuple


class BEVGenerator:
    @staticmethod
    def create_bev(
        point_data: Dict[str, List[Point3D]],
        cfg: Dict[str, Any],
        id: str,
        images: List[str],
        output_folder: str,
        xyrange : Tuple[Tuple, Tuple] = None,
        save_meta : bool = False,
    ) -> None:
        re = cfg.get('resolution', 10.0)

        all_points = []
        for points in point_data.values():
            all_points.extend(points)

        if len(all_points) == 0:
            return
        X = np.array([p.x for p in all_points])
        Y = np.array([p.y for p in all_points])
        Z = np.array([p.z for p in all_points])
        S = np.array([p.strength for p in all_points])
        C = np.array([p.cohen_factor for p in all_points])

        if xyrange is None:
            x_min = np.floor(X.min() * re) / re
            x_max = np.ceil(X.max() * re) / re
            y_min = np.floor(Y.min() * re) / re
            y_max = np.ceil(Y.max() * re) / re
        else:
            x_min = np.floor(xyrange[0][0] * re) / re
            x_max = np.ceil(xyrange[0][1] * re) / re
            y_min = np.floor(xyrange[1][0] * re) / re
            y_max = np.ceil(xyrange[1][1] * re) / re

        width_px = int((x_max - x_min) * re)
        height_px = int((y_max - y_min) * re)

        px = np.floor((X - x_min) * re).astype(int)
        py = np.floor((Y - y_min) * re).astype(int)
        px = np.clip(px, 0, width_px - 1)
        py = np.clip(py, 0, height_px - 1)

        metadata = {
            "left_bottom": [float(x_min), float(y_min)],
            "right_top": [float(x_max), float(y_max)],
            "resolution": re,
            "width_px": width_px,
            "height_px": height_px,
            "lat0_lon0": [cfg.get('lat0'), cfg.get('lon0')],
            "h0": cfg.get('h0'),
            "heading": cfg.get('heading'),
            "start_id": next(iter(point_data.items()))[0],
            "end_id": next(reversed(point_data.items()))[0],
        }

        output_folder = os.path.join(output_folder, cfg.get("data_name"))
        os.makedirs(output_folder, exist_ok=True)
        if save_meta:
            with open(os.path.join(output_folder, f"{id}_meta.json"), 'w') as f:
                json.dump(metadata, f, indent=4)
        
        np.random.seed(42)
        for image_name in images:
            if image_name == images[0]:
                mask = np.ones_like(Z, dtype=bool)
            elif image_name == images[1]: 
                mask = (Z >= 0) & (Z <= 2)
            elif image_name == images[2]:
                mask = (Z >= 0) & (Z <= 0.5)
            elif image_name == images[3]:
                mask = (Z >= -1) & (Z <= 0.5)
            elif image_name == images[4]:
                mask = (np.random.permutation(len(X)) < len(X) // 2)
            elif image_name == images[5]:
                mask = (S > 400000)

            filter_X = X[mask]
            filter_px = px[mask]
            filter_py = py[mask]
            filter_Z = Z[mask]
            filter_S = S[mask]
            filter_C = C[mask]

            rgb = BEVGenerator._compute_rgb(
                filter_Z, filter_S, filter_C
            )
            rgb_sum = np.zeros((height_px, width_px, 3), dtype=np.float32)
            count = np.zeros((height_px, width_px), dtype=int)
            for i in range(len(filter_X)):
                rgb_sum[filter_py[i], filter_px[i]] += rgb[i] * 255
                count[filter_py[i], filter_px[i]] += 1

            # radius = 1
            # for i in range(len(filter_X)):
            #     y, x = filter_py[i], filter_px[i]
            #     for dy in range(-radius, radius+1):
            #         for dx in range(-radius, radius+1):
            #             yy, xx = y + dy, x + dx
            #             if 0 <= yy < height_px and 0 <= xx < width_px:
            #                 rgb_sum[yy, xx] += rgb[i] * 255
            #                 count[yy, xx] += 1
            
            rgb_avg = np.zeros_like(rgb_sum)
            mask = count > 0
            rgb_avg[mask] = rgb_sum[mask] / count[mask][:, None]

            rgb_image = np.flipud(rgb_avg)
            rgb_image = np.clip(rgb_image, 0, 255).astype(np.uint8)
            image_path = os.path.join(output_folder, f"{id}_{image_name}.png")
            Image.fromarray(rgb_image, 'RGB').save(image_path)
    
    @staticmethod
    def _compute_rgb(z: np.ndarray, s: np.ndarray, c: np.ndarray) -> np.ndarray:
        z = BEVGenerator.uniform(z)
        s = BEVGenerator.uniform(s)
        c = BEVGenerator.uniform(c)

        hue = 2/3 * (1 - z)
        value = 1 - s
        saturation = c / (c + 1e-3) * 0.6
        
        hsv = np.stack((hue, saturation, value), axis=1)
        rgb = mcolors.hsv_to_rgb(hsv)
        return rgb

    @staticmethod
    def uniform(data: np.ndarray) -> np.ndarray:
        if len(data) == 0:
            return data
        
        sorted_unique = np.unique(data)
        
        if len(sorted_unique) == 1:
            return np.zeros_like(data)
        
        rank = np.searchsorted(sorted_unique, data)
        normalized = rank / (len(sorted_unique) - 1)
        
        return normalized.reshape(data.shape)
    
    @staticmethod
    def create_combine_bev(
        cfg: Dict[str, Any],
        opacity: float = 0.8
    ) -> None:
        output_folder = cfg.get('combine_output_folder')
        tensor_folder = cfg.get('input_tensor_output_folder')
        label_folder = cfg.get('label_image_output_folder')
        os.makedirs(output_folder, exist_ok=True)
    
        label_files = {f for f in os.listdir(label_folder) 
                      if f.lower().endswith('.png') and '_' in f}
        
        processed_pairs = set()

        for label_name in label_files:
            try:
                base_prefix = label_name.rsplit('_', 1)[0]
                
                tensor_subfolder = '_'.join(base_prefix.split('_')[:2])
                tensor_subfolder_path = os.path.join(tensor_folder, tensor_subfolder)
                all_file_name = f"{base_prefix}_all.png"
                all_file_path = os.path.join(tensor_subfolder_path, all_file_name)
                
                if not os.path.exists(all_file_path):
                    print(f"All file not found: {all_file_path}")
                    continue
                
                if base_prefix in processed_pairs:
                    continue
                
                matching_label_files = [f for f in os.listdir(label_folder) 
                                      if f.startswith(base_prefix + '_') and f.lower().endswith('.png')]
                
                img1 = Image.open(all_file_path).convert('RGBA')
                
                for label_file in matching_label_files:
                    img2_path = os.path.join(label_folder, label_file)
                    img2 = Image.open(img2_path).convert('RGBA')
                    
                    combined = Image.blend(img1, img2, 1 - opacity)
                    
                    parts = label_file.rsplit('_', 1)
                    output_name = f"{parts[1].split('.')[0]}_{parts[0]}.png"
                    output_path = os.path.join(output_folder, output_name)
                    
                    combined.save(output_path)
                
                processed_pairs.add(base_prefix)
                
            except Exception as e:
                print(f"{label_name} combine failed: {str(e)}")
